import torch.utils.data as data
import os
import numpy as np
import random
import matplotlib.pyplot as plt
from data_utilities import *

def get_data(dataset_dir, dataset_name, train_frac, eval_frac):

    dataset = load_compact_pkl_dataset(dataset_dir, dataset_name)

    tx_list = dataset['tx_list']
    rx_list = dataset['rx_list']
    rx_list = rx_list[0:1]
    capture_date_list = dataset['capture_date_list']
    capture_date = capture_date_list

    train_dataset = merge_compact_dataset(dataset, capture_date, tx_list, rx_list, equalized=0)
    train_augset,eval_augset,_ = prepare_dataset(train_dataset, tx_list, val_frac=eval_frac, test_frac=1.-train_frac-eval_frac)
    [sig_train,txid_train,_,_] = train_augset
    [sig_valid,txid_valid,_,_] = eval_augset
    train=sig_valid
    txid=txid_valid
    sig_train=train[0:8000]#改成8000训练集，3200测试集。源代码是224测试集，11200训练集
    sig_valid=train[8000:11200]
    txid_train=txid[0:8000]
    txid_valid=txid[8000:11200]
    return sig_train, txid_train, sig_valid, txid_valid

class SigLoader(data.Dataset):
    def __init__(self, data, label, transform=None):
        self.data = np.transpose(data, [0,2,1])
        self.label = np.array(label, dtype=np.int64)
        print('number of samples:', len(self.label))
        self.transform = transform

    def __getitem__(self, index: int):
        x = self.data[index,:,:]
        y = self.label[index]
        if self.transform is not None:
            x = self.transform(x)
        return x, y
    
    def __len__(self):
        return len(self.label)